import torch
from torch import nn
from torch.nn import functional as F


IMG_SIZE = 224
LR = 0.005


class DenseBlock(nn.Module):
    def __init__(self, in_channels, growth_rate, num_layers, bn_size=4):
        super(DenseBlock, self).__init__()
        self.layers = nn.Sequential()
        for _ in range(num_layers):
            self.layers.append(nn.Sequential(
                nn.BatchNorm2d(in_channels),
                nn.ReLU(),
                nn.Conv2d(in_channels, bn_size*growth_rate, kernel_size=1),
                nn.BatchNorm2d(bn_size*growth_rate),
                nn.ReLU(),
                nn.Conv2d(bn_size*growth_rate, growth_rate, kernel_size=3, padding=1),
            ))
            in_channels += growth_rate

    def forward(self, input):
        for net in self.layers:
            y = net(input)
            y = torch.concat((y, input), dim=1)
            input = y
        return y


class TransitionBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(TransitionBlock, self).__init__()
        self.net = nn.Sequential(
            nn.BatchNorm2d(in_channels),
            nn.ReLU(),
            nn.Conv2d(in_channels, out_channels, kernel_size=1),
            nn.AvgPool2d(kernel_size=2, stride=2)
        )

    def forward(self, input):
        return self.net(input)


class DenseNet121(nn.Module):
    """DenseNet-121"""
    def __init__(self, num_classes, in_channels=3):
        super(DenseNet121, self).__init__()
        self.features = nn.Sequential(
            ## Stage0 7x7 conv, stride 2
            nn.Conv2d(in_channels, 64, kernel_size=7, stride=2, padding=3),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),

            ## Stage1
            DenseBlock(64, 32, 6),
            TransitionBlock(256, 128),

            ## Stage2
            DenseBlock(128, 32, 12),
            TransitionBlock(512, 256),

            ## Stage3
            DenseBlock(256, 32, 24),
            TransitionBlock(1024, 512),

            ## Stage4
            DenseBlock(512, 32, 16),
            TransitionBlock(1024, 1024),
        )
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(1024, num_classes)
        )


    def forward(self, input):
        y = self.features(input)
        y = self.classifier(y)
        return y
    

    @property
    def image_size(self):
        return IMG_SIZE
    

    @property
    def lr(self):
        return LR


class DenseNet169(nn.Module):
    """DenseNet-169"""
    def __init__(self, num_classes, in_channels=3):
        super(DenseNet169, self).__init__()
        self.features = nn.Sequential(
            ## Stage0 7x7 conv, stride 2
            nn.Conv2d(in_channels, 64, kernel_size=7, stride=2, padding=3),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),

            ## Stage1
            DenseBlock(64, 32, 6),
            TransitionBlock(256, 128),

            ## Stage2
            DenseBlock(128, 32, 12),
            TransitionBlock(512, 256),

            ## Stage3
            DenseBlock(256, 32, 32),
            TransitionBlock(1280, 512),

            ## Stage4
            DenseBlock(512, 32, 32),
            TransitionBlock(1536, 1024),
        )
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(1024, num_classes)
        )


    def forward(self, input):
        y = self.features(input)
        y = self.classifier(y)
        return y
    

    @property
    def image_size(self):
        return IMG_SIZE
    

    @property
    def lr(self):
        return LR


class DenseNet201(nn.Module):
    """DenseNet-201"""
    def __init__(self, num_classes, in_channels=3):
        super(DenseNet201, self).__init__()
        self.features = nn.Sequential(
            ## Stage0 7x7 conv, stride 2
            nn.Conv2d(in_channels, 64, kernel_size=7, stride=2, padding=3),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),

            ## Stage1
            DenseBlock(64, 32, 6),
            TransitionBlock(256, 128),

            ## Stage2
            DenseBlock(128, 32, 12),
            TransitionBlock(512, 256),

            ## Stage3
            DenseBlock(256, 32, 48),
            TransitionBlock(1792, 512),

            ## Stage4
            DenseBlock(512, 32, 32),
            TransitionBlock(1536, 1024),
        )
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(1024, num_classes)
        )


    def forward(self, input):
        y = self.features(input)
        y = self.classifier(y)
        return y
    

    @property
    def image_size(self):
        return IMG_SIZE
    

    @property
    def lr(self):
        return LR
    

class DenseNet264(nn.Module):
    """DenseNet-264"""
    def __init__(self, num_classes, in_channels=3):
        super(DenseNet264, self).__init__()
        self.features = nn.Sequential(
            ## Stage0 7x7 conv, stride 2
            nn.Conv2d(in_channels, 64, kernel_size=7, stride=2, padding=3),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),

            ## Stage1
            DenseBlock(64, 32, 6),
            TransitionBlock(256, 128),

            ## Stage2
            DenseBlock(128, 32, 12),
            TransitionBlock(512, 256),

            ## Stage3
            DenseBlock(256, 32, 64),
            TransitionBlock(32*64+256, 512),

            ## Stage4
            DenseBlock(512, 32, 48),
            TransitionBlock(32*48+512, 1024),
        )
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(1024, num_classes)
        )


    def forward(self, input):
        y = self.features(input)
        y = self.classifier(y)
        return y
    

    @property
    def image_size(self):
        return IMG_SIZE
    

    @property
    def lr(self):
        return LR